import numpy as np
from sklearn.model_selection import train_test_split

class TrainingPipeline:
    def __init__(self, data_path):
        self.data = np.load(data_path)
        self.X = self.data['features']
        self.y = self.data['labels']
        
    def run(self, model_type='random_forest'):
        """完整训练流程"""
        X_train, X_test, y_train, y_test = train_test_split(
            self.X, self.y, test_size=0.2, stratify=self.y
        )
        
        if model_type in ['random_forest', 'svm']:
            trainer = TraditionalModelTrainer(model_type)
            model = trainer._train(X_train, y_train)
            cv_scores = trainer.cross_validate(X_train, y_train)
        else:
            trainer = DeepLearningTrainer(
                input_shape=(X_train.shape[1], 1),
                model_type=model_type
            )
            model = trainer.build_model()
            model.fit(X_train, y_train, epochs=20, batch_size=32)
            cv_scores = trainer.cross_validate(X_train, y_train)
            
        test_acc = self._evaluate(model, X_test, y_test)
        return {
            'model': model,
            'cv_scores': cv_scores,
            'test_accuracy': test_acc
        }
    
    def _evaluate(self, model, X_test, y_test):
        """测试集评估"""
        if hasattr(model, 'predict_proba'):
            y_pred = model.predict(X_test).round()
        else:
            y_pred = (model.predict(X_test) > 0.5).astype(int)
        return accuracy_score(y_test, y_pred) 